import datetime
import logging
import os
import time
import traceback
from absl import app, flags
from collections import namedtuple
import pandas as pd
import numpy as np
import matplotlib
import matplotlib.pyplot as plt

import coin.proto.coin_request_pb2 as cr_pb2
import coin2.service.strategy.trading_summary_service_pb2 as tss_pb2
import coin2.service.account.account_info_service_pb2 as ais_pb2
from xunkemgmt_client.base.datetime_util import to_timestamp_int, to_datetime
from xunkemgmt_client.client.api_client import XunkemgmtClient
from python.experimental.lucasyoo.accounting.offset_accounting2 import (Data)
import python.experimental.lucasyoo.accounting.accounting as lac

matplotlib.use("Agg")
os.environ["TZ"] = "GMT"
time.tzset()

FLAGS = flags.FLAGS

STRAT_TRADING_PNL = namedtuple('STRAT_PNL', ['strategy_name', 'trading_date', 'trading_pnl_in_usd'])


def plot_and_save(df: pd.DataFrame, strat_name, output_dir: str):
  start_dt = df.index.min().strftime("%Y%m%d")
  end_dt = df.index.max().strftime("%Y%m%d")
  df.to_csv(f'{output_dir}/{strat_name}_{start_dt}_{end_dt}.csv')
  plt.rcParams['figure.figsize'] = (30, 20)
  plt.plot(df.index, df['pnl_in_usd'], 'k-', lw=1.0, drawstyle='steps-post', label='pnl_in_usd')
  plt.plot(df.index,
           df['trading_pnl_in_usd'],
           'r-',
           lw=0.5,
           drawstyle='steps-post',
           label='trading_pnl_in_usd')
  plt.plot(df.index,
           df['otc_pnl_in_usd'],
           'b-',
           lw=0.5,
           drawstyle='steps-post',
           label='otc_pnl_in_usd')
  plt.legend()
  plt.title(f"{strat_name} ({start_dt}-{end_dt})"
            f"(total: {df['pnl_in_usd'].iloc[-1]:,.0f},"
            f" trading: {df['trading_pnl_in_usd'].iloc[-1]:,.0f},"
            f" option: {df['otc_pnl_in_usd'].iloc[-1]:,.0f})")
  plt.savefig(f'{output_dir}/{strat_name}_{start_dt}_{end_dt}.png')
  plt.close()


def plot_total(start_time: datetime.datetime, end_time: datetime.datetime, output_dir: str):
  res = []
  strat_base = []
  tmp_dfs = []
  kline_dfs = []
  hypo_reserve_infos = {}
  with XunkemgmtClient() as client:
    rsp = client.query_trading_summary_interval_histories(
        tss_pb2.QueryTradingSummaryIntervalHistoriesRequestProto(
            business_units=['Day1mm'],
            strategy_groups=['vmm', 'hamm'],
            start_timestamp=to_timestamp_int(start_time - datetime.timedelta(days=1)),
            end_timestamp=to_timestamp_int(end_time),
            agg_types=['STRATEGY_NAME', 'TRADING_DATE']))
    for his in rsp.histories:
      strategy_name = his.summary_info.strategy_name
      trading_date = his.summary_info.trading_date
      trading_pnl_in_usd = his.summary.trading_pnl_mark_in_usd
      base = strategy_name.replace("vmm_", "").replace("hamm_", "").replace("_delta_hedge", "")
      if "_" in base:
        base = base.split("_")[0]
      base = base.upper()
      if FLAGS.symbol != "ALL" and (base not in FLAGS.symbol.split(",")):
        continue
      tmp = STRAT_TRADING_PNL(strategy_name, str(trading_date), trading_pnl_in_usd)
      res.append(tmp)
      strat_base.append((strategy_name, base, str(trading_date)))
  strat_info = pd.DataFrame(strat_base, columns=["strategy_name", "base", "trading_date"])
  strat_info['trading_date'] = pd.to_datetime(strat_info['trading_date'])
  unique_strat_info = strat_info[['strategy_name', 'base']].drop_duplicates()
  for _, (strategy_name, base) in unique_strat_info.iterrows():
    tmp_start_time = strat_info.loc[(strat_info['strategy_name'] == strategy_name) &
                                    (strat_info['base'] == base),
                                    'trading_date'].min() - datetime.timedelta(days=1)
    # tmp_end_time = strat_info.loc[(strat_info['strategy_name'] == strategy_name) &
    #                               (strat_info['base'] == base), 'trading_date'].max()
    if "delta_hedge" not in strategy_name:
      continue
    tmp_df, kline_df, hypo_reserve_info = lac.get_option_pnl(tmp_start_time, end_time,
                                                             strategy_name, base)
    tmp_df['strategy_name'] = strategy_name
    tmp_df['base'] = base
    tmp_dfs.append(tmp_df)
    kline_dfs.append(kline_df)
    hypo_reserve_infos.update(hypo_reserve_info)
  option_df = pd.concat(tmp_dfs)
  if not option_df.empty:
    sliced_df = option_df[['ts', 'base', 'strategy_name', 'call_price_diff']].copy()
    config = lac.query_option_num_info()
    sliced_df = sliced_df.merge(pd.DataFrame(config).T.reset_index(names="base"),
                                on="base",
                                how='left')
    sliced_df.rename(columns={'ts': 'trading_date'}, inplace=True)
    sliced_df['otc_pnl_in_usd'] = sliced_df['call_price_diff'] * sliced_df['team_flow']
  else:
    sliced_df = pd.DataFrame(columns=['trading_date', 'base', 'strategy_name', 'otc_pnl_in_usd'])
  df = pd.DataFrame(res)
  df['trading_date'] = pd.to_datetime(df['trading_date'])
  kline_dfs = pd.concat(kline_dfs)
  if not kline_dfs.empty:
    kline_dfs[["market_type", "exchange",
               "symbol"]] = kline_dfs['product_str'].str.split(':', expand=True)
    kline_dfs[['base', 'quote']] = kline_dfs['symbol'].str.split('-', expand=True)
    kline_dfs['trading_date'] = pd.to_datetime(kline_dfs['klineTimestamp'])
  for strategy_name in df['strategy_name'].unique():
    if strategy_name.startswith("hamm_") or strategy_name in [
        'vmm_joy_usdt', 'vmm_fncy_usdt', 'vmm_cube_usdt'
    ]:
      continue
    tmp_df = df[df['strategy_name'] == strategy_name]
    if tmp_df.empty:
      continue
    if tmp_df['trading_date'].max() + datetime.timedelta(days=1) < end_time:
      tmp_bal = []
      with XunkemgmtClient() as client:
        rsp1 = client.query_applications(cr_pb2.ApplicationRequestProto())
        strat_acc_id_map = {
            i.account_id: i.strategy_name for i in rsp1 if i.strategy_name == strategy_name
        }
        account_ids = list(strat_acc_id_map.keys())
        rsp = client.query_balance_interval_histories(
            ais_pb2.QueryBalanceIntervalHistoriesRequestProto(
                account_ids=account_ids,
                start_timestamp=to_timestamp_int(tmp_df['trading_date'].max() -
                                                 datetime.timedelta(days=1)),
                end_timestamp=to_timestamp_int(end_time)))
        for bal_hist in rsp.histories:
          trading_date = bal_hist.balance_info.trading_date
          acc_id = bal_hist.balance_info.account_id
          for each_bal in bal_hist.each_balance:
            if each_bal.currency in unique_strat_info['base'].unique():
              tmp_bal.append((datetime.datetime.strptime(str(trading_date), "%Y%m%d") +
                              datetime.timedelta(days=1), f"{bal_hist.balance_info.market_type}",
                              f"{bal_hist.balance_info.exchange}", f"{each_bal.currency}",
                              each_bal.total, strat_acc_id_map[acc_id]))
      tmp_bal_df = pd.DataFrame(
          tmp_bal,
          columns=["trading_date", "market_type", "exchange", "base", "total", "strategy_name"])
      tmp_bal_df['strat_prod'] = tmp_bal_df['strategy_name'] + "|" + tmp_bal_df["base"]
      tmp_bal_df["hypo_reserve"] = tmp_bal_df['strat_prod'].apply(
          lambda x: hypo_reserve_infos.get(x))
      add_df = pd.merge(
          tmp_bal_df[["trading_date", "base", "exchange", "total", "strategy_name",
                      "hypo_reserve"]],
          kline_dfs[["trading_date", "base", "exchange", "open", "close"]],
          on=["trading_date", "base", "exchange"],
          how="left")
      add_df["trading_pnl_in_usd"] = (add_df["close"] - add_df["open"]) * (add_df["total"] -
                                                                           add_df["hypo_reserve"])
      add_trd_pnl_df = add_df[["strategy_name", "trading_date", "trading_pnl_in_usd"]]
      tmp_df2 = tmp_df[tmp_df['trading_date'].isin(set(tmp_df['trading_date']) - set(add_trd_pnl_df['trading_date']))]
      tmp_df2 = pd.concat([tmp_df2, add_trd_pnl_df])
      df = df.loc[(df['strategy_name'] != strategy_name)]
      df = pd.concat([df, tmp_df2])
  final_df = pd.merge(df,
                      sliced_df[["trading_date", "strategy_name", "otc_pnl_in_usd"]],
                      on=['strategy_name', 'trading_date'],
                      how='outer')
  final_df.fillna(0.0, inplace=True)
  final_df['pnl_in_usd'] = final_df['trading_pnl_in_usd'].add(final_df['otc_pnl_in_usd'],
                                                              fill_value=0)
  total_pnl = final_df.groupby(['trading_date'
                               ])[['pnl_in_usd', 'trading_pnl_in_usd',
                                   'otc_pnl_in_usd']].sum().cumsum()
  plot_and_save(total_pnl, 'TOTAL', output_dir)
  for strat in final_df["strategy_name"].unique():
    tmp = final_df[final_df["strategy_name"] == strat]
    tmp = tmp.groupby(['trading_date'])[['pnl_in_usd', 'trading_pnl_in_usd',
                                         'otc_pnl_in_usd']].sum().cumsum()
    plot_and_save(tmp, strat, output_dir)


def main(_):
  start_time = datetime.datetime.strptime(FLAGS.start_time, '%Y%m%d')
  end_time = datetime.datetime.strptime(FLAGS.end_time, '%Y%m%d')
  output_dir = os.path.join(os.path.expandvars('$COIN_REPO'), FLAGS.output_dir)
  os.makedirs(output_dir, exist_ok=True)
  plot_total(start_time, end_time, output_dir)


if __name__ == "__main__":
  flags.DEFINE_string('start_time', '20230101', '%Y%m%d')
  flags.DEFINE_string('end_time', '20230331', '%Y%m%d')
  flags.DEFINE_string('output_dir', 'sim_result/accounting', 'output dir')
  flags.DEFINE_string('symbol', 'ALL', 'symbol')
  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
